import tensorflow as tf
import logging

class CustomCallback(tf.keras.callbacks.Callback):
    """自定义训练回调函数"""
    
    def __init__(self):
        super(CustomCallback, self).__init__()
        self.logger = logging.getLogger(__name__)
    
    def on_epoch_begin(self, epoch, logs=None):
        self.logger.info(f"开始第 {epoch + 1} 轮训练")
    
    def on_epoch_end(self, epoch, logs=None):
        if logs:
            self.logger.info(
                f"第 {epoch + 1} 轮结束 - "
                f"损失: {logs.get('loss', 0):.4f}, "
                f"验证损失: {logs.get('val_loss', 0):.4f}"
            )
    
    def on_train_begin(self, logs=None):
        self.logger.info("训练开始")
    
    def on_train_end(self, logs=None):
        self.logger.info("训练结束")